from fastapi import APIRouter, File, UploadFile, HTTPException, Request
from fastapi.responses import JSONResponse
from pydantic import BaseModel
import joblib
import torch
from PIL import Image
from torchvision import transforms, models
import torch.nn as nn
import numpy as np
import io
import os

router = APIRouter(prefix="", tags=["Predict"])

# ----------------------------
# 初始化特征提取模型
# ----------------------------
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

feature_extractor = models.inception_v3(pretrained=True)
feature_extractor.fc = nn.Identity()
feature_extractor.eval()
feature_extractor.to(device)

transform = transforms.Compose([
    transforms.Resize(299),
    transforms.CenterCrop(299),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225])
])

# 模型缓存
model_cache = {}


class PredictParams(BaseModel):
    model_path: str
    class_list: list[str]


@router.post("/predict")
async def predict_dog_breed(
    request: Request,
    file: UploadFile = File(...)
):
    """
    上传图片文件 + JSON body
    JSON body 示例：
    {
      "model_path": "models/svm_model_93.pkl",
      "class_list": ["golden_retriever", "poodle", "husky"]
    }

    curl 示例：
    curl -X POST "http://127.0.0.1:8000/predict" \
      -F "file=@dog.jpg" \
      -F 'json={"model_path":"models/svm_model_93.pkl","class_list":["golden_retriever","poodle","husky"]}'
    """
    try:
        # 读取 JSON body
        form = await request.form()
        json_data = form.get("json")
        if not json_data:
            raise HTTPException(status_code=400, detail="缺少 json 字段")

        import json
        try:
            params = PredictParams(**json.loads(json_data))
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"JSON 解析失败: {str(e)}")

        # 检查模型路径
        if not os.path.exists(params.model_path):
            raise HTTPException(status_code=400, detail=f"模型文件不存在: {params.model_path}")

        if not file.content_type.startswith("image/"):
            raise HTTPException(status_code=400, detail="上传文件必须是图片")

        # 1️⃣ 读取图片
        contents = await file.read()
        img = Image.open(io.BytesIO(contents)).convert("RGB")
        img_tensor = transform(img).unsqueeze(0).to(device)

        # 2️⃣ 提取特征
        with torch.no_grad():
            features = feature_extractor(img_tensor)
        features_np = features.cpu().numpy()

        # 3️⃣ 加载模型（使用缓存）
        if params.model_path not in model_cache:
            model_cache[params.model_path] = joblib.load(params.model_path)
        svm_model = model_cache[params.model_path]

        # 4️⃣ 预测
        pred_label = svm_model.predict(features_np)[0]
        if pred_label >= len(params.class_list):
            breed = f"未知类别({pred_label})"
        else:
            breed = params.class_list[pred_label]

        return JSONResponse(content={
            "predicted_breed": breed,
            "pred_label_index": int(pred_label),
            "model_used": params.model_path,
            "device": str(device)
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"预测失败: {str(e)}")
